package cz.cuni.amis.episodic.lisp.netcreators;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import jsmile.utils.CsvToSmileData;
import lombok.Cleanup;

import org.apache.commons.lang.StringUtils;

import smile.Network;
import cz.cuni.amis.episodic.lisp.LispPlanToCsvData;
import cz.cuni.amis.episodic.lisp.LispPlanToCsvData.LispToCsvResult;
import cz.cuni.amis.episodic.lisp.behan.LispTree;
import cz.cuni.amis.episodic.lisp.behan.TraceToEventVisitor;
import cz.cuni.amis.episodic.lisp.visitor.HHMMmpTraceVisitor;

/**
 * Sequence of altering Goal and Hidden nodes. Hidden nodes are learned by EM.
 * Finish nodes are needed to achieve power of hierarchical FSM.
 * 
 * @author ik
 */
public class AHMEMCreator extends NetCreator {
	/**
	 * Maximum number of episode levels extracted from the data.
	 */
	int maxDepth = 2;

	/**
	 * Number of states of the memory node.
	 */
	int hiddenNodeStatesNum;
	
	boolean observationTemporalModel;

	public AHMEMCreator(int hiddenNodeStatesNum, int levels) {
		this(hiddenNodeStatesNum, levels, false);
	}
	
	public AHMEMCreator(int hiddenNodeStatesNum, int levels, boolean observationTemporalModel) {
		this.hiddenNodeStatesNum = hiddenNodeStatesNum;
		this.observationTemporalModel = observationTemporalModel;
		String observ = observationTemporalModel ? "-obs" : "";
		filenameSuffix = "_ahmem" + observ + "^" + levels + "_"
				+ hiddenNodeStatesNum;
		maxDepth = levels;
	}

	public AHMEMCreator(int hiddenNodeStatesNum,
			TraceModificationStrategy strategy, boolean observationTemporalModel) {
		this.hiddenNodeStatesNum = hiddenNodeStatesNum;
		this.observationTemporalModel = observationTemporalModel;
		String observ = observationTemporalModel ? "-obs" : "";
		filenameSuffix = "_ahmem" + observ + "^" + strategy.getDepth() + "_"
				+ hiddenNodeStatesNum;
		setTraceModificationStrategy(strategy);
		// maxDepth = levels;
	}

	public AHMEMCreator(int hiddenNodeStatesNum,
			TraceModificationStrategy strategy) {
		this(hiddenNodeStatesNum, strategy, false);
	}
	
	public List<List<String>> lispToTrainingData(File inputLispPlan,
			File smileTrainingData) throws IOException {
		// get filename without path and extension

		// read lisp to csv
		ByteArrayOutputStream csvOs = new ByteArrayOutputStream();
		LispPlanToCsvData.LispToCsvResult res = lispToCsv(
				inputLispPlan.getAbsolutePath(), csvOs, batchSize,
				traceModificationStrategy);
		csvOs.close();

		// write also the csv
		String smilePath = smileTrainingData.getAbsolutePath();
		File csvFile = new File(smilePath.substring(0,
				smilePath.lastIndexOf('.') + 1)
				+ "csv");
		@Cleanup
		FileOutputStream fos = new FileOutputStream(csvFile);
		fos.write(csvOs.toByteArray());

		// transform csv to smile data format
		CsvToSmileData.convert(new InputStreamReader(new ByteArrayInputStream(
				csvOs.toByteArray())), smileTrainingData.getAbsolutePath());

		return res.levelActions;
	}

	/**
	 * UGLY cut and paste and edit from LispPlanToCsvData ... should be
	 * refactored.
	 * 
	 * @param inputLispPlan
	 * @param outputPlanCsv
	 * @param batchSize
	 * @param translator
	 * @return
	 * @throws FileNotFoundException
	 */
	public static LispPlanToCsvData.LispToCsvResult lispToCsv(
			String inputLispPlan, OutputStream outputPlanCsv, int batchSize,
			TraceModificationStrategy strategy) throws FileNotFoundException {
		LispTree parser = new LispTree(inputLispPlan);
		LispPlanToCsvData.LispToCsvResult res = new LispPlanToCsvData.LispToCsvResult();

		HHMMmpTraceVisitor visitor = new HHMMmpTraceVisitor();
		parser.walkInForest(visitor);

		// TODO discard higher levels of hierarchy
		// visitor.trimHighest(maxDepth);
		// visitor.paddingUpper(maxDepth);

		visitor.modify(strategy);

		// find max depth
		int max = Integer.MIN_VALUE;
		for (String[] trace : visitor.getTraces()) {
			max = Math.max(max, trace.length);
		}

		@Cleanup
		PrintStream out = new PrintStream(outputPlanCsv);
		String delim = "\t";

		// each set contains actions found on the corresponding level in
		// hierarchy
		Set<String>[] levels = new HashSet[max];
		for (int i = 0; i < max; i++) {
			levels[i] = new HashSet<String>();
		}

		// print header
		String[] nodesInHeader = new String[max * 2 - 1];
		for (int i = 0; i < nodesInHeader.length - 1; i += 2) {
			int layerIx = max - (i / 2 + 1);
			nodesInHeader[i] = "G" + layerIx;
			nodesInHeader[i + 1] = "F" + layerIx;
		}
		nodesInHeader[nodesInHeader.length - 1] = "O";
		List<String> header = new ArrayList<>(Arrays.asList("Batch", "Id"));
		header.addAll(Arrays.asList(nodesInHeader));
		out.println(StringUtils.join(header, delim));

		
		// output with padding

		int batchNum = 1;
		int linesInBatch = 0;
		String[] line = new String[max];
		int traceIx = 0;
		int finishTraceIx = 0;
		Boolean[] finishPadding = new Boolean[visitor.getTraces().get(0).length];
		Arrays.fill(finishPadding, true);
		
		for (String[] trace : visitor.getTraces()) {
			linesInBatch++;
			out.print(batchNum + delim + Integer.toString(res.lineNum++)
					+ delim);

			// join goal nodes and finish nodes
			String[] outLine = new String[trace.length * 2 - 1];
			
			// ad hoc fix for finish vals when day padding is enabled
			if(!trace[0].equals("DAY_END")) {
				joinToSingleLine(outLine, trace, visitor.getFinish().get(finishTraceIx)); 
				finishTraceIx++;
			} else {
				joinToSingleLine(outLine, trace, finishPadding);
			}
			
			
			outLine[outLine.length - 1] = trace[trace.length - 1];

			out.println(StringUtils.join(outLine, delim));

			if (batchSize > 0 && linesInBatch == batchSize) {
				batchNum++;
				linesInBatch = 0;
			}

			// remember actions/goals found in each level
			for (int i = 0; i < trace.length; i++) {
				levels[i].add(trace[i]);
			}
			traceIx++;
		}

		// create result
		res.batchSize = batchSize;
		res.inputLispPlan = inputLispPlan;

		// sort actions by alphabet
		for (int i = 0; i < levels.length; i++) {
			List<String> actions = new ArrayList<>(levels[i]);
			Collections.sort(actions);
			res.levelActions.add(actions);
		}
		return res;
	}

	static void joinToSingleLine(String[] outLine, String[] trace, Boolean[] finish) {
		for (int i = 0; i < trace.length - 1; i++) {
			outLine[2 * i] = trace[i];
			outLine[2 * i + 1] = finish[i] ? "finish"
					: "continue";
		}
	}
	
	/**
	 * Add finish states.
	 * 
	 * @return
	 */
	static boolean[][] transformToHHMMFinish(List<List<String>> data) {
		boolean[][] finish = new boolean[data.size()][];
		List<String> lastTrace = null;
		int traceIx = 0;
		for (List<String> trace : data) {
			if (lastTrace != null) {
				boolean[] finishData = new boolean[trace.size()];
				// compare
				for (int i = 0; i < trace.size(); i++) {
					finishData[i] = !lastTrace.get(i).equals(trace.get(i));
				}
				finish[traceIx] = finishData;
				traceIx++;
			} else {
				// first trace, nothing to compare
			}
			lastTrace = trace;
		}
		// on last slice every state finishes
		boolean[] lastFinish = new boolean[lastTrace.size()];
		Arrays.fill(lastFinish, true);
		finish[finish.length - 1] = lastFinish;
		return finish;
	}

	@Override
	public Network createNet(List<List<String>> levelActions) {
		List<String> statesF = Arrays.asList("continue", "finish");

		final int max = levelActions.size();

		Network net = new Network();
		String[] nodes = createNodeNames(max);

		// observation O
		addNode(net, nodes[nodes.length - 1],
				levelActions.get(levelActions.size() - 1));
		// add rest if the nodes
		for (int i = 0; i < max - 1; i++) {
			int id = max - i - 1;
			int nodeIx = 2 * i;
			String finishNode = "F" + id;

			// F ... finish node
			addNode(net, finishNode, statesF);
			deleteDefaultOutcomes(net, finishNode);

			// G
			addNode(net, nodes[nodeIx], levelActions.get(i));

			// H
			addNode(net, nodes[nodeIx + 1], hiddenNodeStatesNum);
		}

		// create arcs between nodes
		for (int i = 0; i < max - 1; i++) {
			int id = max - i - 1;
			int nodeIx = 2 * i;
			String finishNode = "F" + id;

			if (i > 0) {
				// H_(i-1) -> G_i
				net.addArc(nodes[nodeIx - 1], nodes[nodeIx]);

				// F_i -> F_(i+1)
				net.addArc("F" + i, "F" + (i + 1));
			}
			// G_i -> H_i
			net.addArc(nodes[nodeIx], nodes[nodeIx + 1]);
			// G_i -> G_(i-1)
			net.addArc(nodes[nodeIx], nodes[nodeIx + 2]);
			// G -> F
			net.addArc(nodes[nodeIx], finishNode);
			// H -> F
			net.addArc(nodes[nodeIx + 1], finishNode);
			// F -> G
			net.addTemporalArc(finishNode, nodes[nodeIx], 1);
			// F -> H
			net.addTemporalArc(finishNode, nodes[nodeIx + 1], 1);

			// graphics
			// TODO net.setNodePosition(finishNode, x, y, width, height);
		}
		// H_last -> O
		net.addArc(nodes[nodes.length - 2], nodes[nodes.length - 1]);

		// add temporal arcs
		adjustNodes(net, nodes, observationTemporalModel);

		return net;

	}

	static void adjustNodes(Network net, String[] nodes,
			boolean observationTemporalModel) {
		// values for layout
		int x = 180;
		int y = 60;
		int width = 80;
		int height = 60;
		for (String node : nodes) {
			// net.setNodeTemporalType(node, Network.NodeTemporalType.Plate);

			if (!(observationTemporalModel == false && node.equals("O"))) {
				// net.setNodeTemporalType(node,
				// Network.NodeTemporalType.Plate);
				net.addTemporalArc(node, node, 1);
			}
			// delete default outcomes
			deleteDefaultOutcomes(net, node);

			// net.deleteOutcome(node, 0);
			// net.deleteOutcome(node, 1);

			// do layouting
			net.setNodePosition(node, x, y, width, height);
			y += 80;
		}
	}

	public static void addNode(Network net, String name, int statesNum) {
		addTemporalNode(net, name);
		for (int i = 0; i < statesNum; i++) {
			net.addOutcome(name, "HiddenState" + i);
		}
	}

	public static String[] createNodeNames(int levels) {
		String[] nodes = new String[levels * 2 - 1];
		for (int i = 0; i < levels - 1; i++) {
			int id = levels - i - 1;
			int nodeIx = 2 * i;
			// goal node
			nodes[nodeIx] = "G" + id;
			// hidden node representing goal's internal structure
			nodes[nodeIx + 1] = "H" + id;

		}
		nodes[nodes.length - 1] = "O";
		return nodes;
	}
}
